package com.tsy.pointcloud.sse;

import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class SseService {

    // 存储所有的 SseEmitter 实例
    private final Map<String, SseEmitter> sseEmitters = new ConcurrentHashMap<>();

    /**
     * 创建一个 SseEmitter 实例，并将其存储到 Map 中
     * @param clientId
     * @return
     */
    public SseEmitter createEmitter(String clientId) {
        // 通过设置 SseEmitter 的超时时间来减少超时问题，例如设置为无限期（0L）：
        SseEmitter emitter = new SseEmitter(0L);
        sseEmitters.put(clientId, emitter);

        // 连接断开
        emitter.onCompletion(() -> sseEmitters.remove(clientId));
        // 连接超时
        emitter.onTimeout(() -> sseEmitters.remove(clientId));
        // 连接异常
        emitter.onError(e -> sseEmitters.remove(clientId));

        try {
            // 向客户端发送一条连接成功的事件
            emitter.send(SseEmitter.event().comment("connected"));
        } catch (IOException e) {
            // 如果发送消息失败，则从映射表中移除 emitter
            sseEmitters.remove(clientId);
        }

        return emitter;
    }

    /**
     * 断开指定用户的 SSE 连接
     */
    public void disconnect(String clientId) {
        SseEmitter sseEmitter = sseEmitters.get(clientId);
        if (sseEmitter != null) {
            try {
                sseEmitter.send(SseEmitter.event().comment("disconnected"));
            } catch (Exception ignore) {
            }
            sseEmitters.remove(clientId);
        }
    }

    /**
     * 指定发送消息
     * @param clientId
     * @param message
     * @throws IOException
     */
    public void sendMessage(String clientId, Object message) throws IOException {
        SseEmitter emitter = sseEmitters.get(clientId);
        if (emitter != null) {
            emitter.send(message);
        }
    }

    /**
     * 广播消息
     * @param message
     */
    public void broadcastMessage(Object message) {
        sseEmitters.forEach((clientId, emitter) -> {
            try {
                emitter.send(message);
            } catch (IOException e) {
                emitter.completeWithError(e);
                sseEmitters.remove(clientId);
            }
        });
    }
}